package track2.bayes;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.util.HashMap;
import java.util.Map.Entry;

import track2.context.Config;
import track2.data.*;

/**
 * Training Class for NaiveBayes
 * 
 * @author lingjuan, yijia
 * 
 */
public class TrainClassifier {
	private static HashMap<String, Double> titleC;
	private static HashMap<String, Double> titleNC;
	private static HashMap<String, Double> queryC;
	private static HashMap<String, Double> queryNC;
	private static HashMap<String, Double> descriptionC;
	private static HashMap<String, Double> descriptionNC;
	private static HashMap<String, Double> purchaseC;
	private static HashMap<String, Double> purchaseNC;
	private static HashMap<String, Double> genderC;
	private static HashMap<String, Double> genderNC;
	private static HashMap<String, Double> ageC;
	private static HashMap<String, Double> ageNC;

	private static double titleCCount = 0;
	private static double titleNCCount = 0;
	private static double queryCCount = 0;
	private static double queryNCCount = 0;
	private static double desCCount = 0;
	private static double desNCCount = 0;
	private static double purchaseCCount = 0;
	private static double purchaseNCCount = 0;
	private static double genderCCount = 0;
	private static double genderNCCount = 0;
	private static double ageCCount = 0;
	private static double ageNCCount = 0;
	private static double totalClk = 0;
	private static double totalImp = 0;

	/**
	 * Main Function for training Call this to get a {@link NaiveBayesFeatures}
	 * 
	 * @return NaiveBayesFeatures
	 */
	public static NaiveBayesFeatures train() {
		NaiveBayesFeatures nbFeatures = new NaiveBayesFeatures();

		initial();
		// get the counts for each feature
		trainFeature(Config.trainingFilePath);

		// compute the probability for each feature;
		// Flag "T": title; "Q": query, "D": description, "P":purchaseKeyword,
		// "U": user
		computeProbability(nbFeatures, "T");
		computeProbability(nbFeatures, "Q");
		computeProbability(nbFeatures, "D");
		computeProbability(nbFeatures, "P");
		computeProbability(nbFeatures, "U");
		computeProbability(nbFeatures, "Prior");

		// output all the features to file;

		return nbFeatures;
	}

	private static void initial() {
		titleC = new HashMap<String, Double>();
		titleNC = new HashMap<String, Double>();
		queryC = new HashMap<String, Double>();
		queryNC = new HashMap<String, Double>();
		descriptionC = new HashMap<String, Double>();
		descriptionNC = new HashMap<String, Double>();
		purchaseC = new HashMap<String, Double>();
		purchaseNC = new HashMap<String, Double>();
		genderC = new HashMap<String, Double>();
		genderNC = new HashMap<String, Double>();
		ageC = new HashMap<String, Double>();
		ageNC = new HashMap<String, Double>();
	}

	private static void trainFeature(String trainingFile) {
		try {
			BufferedReader br = new BufferedReader(new FileReader(
					Config.trainingFilePath));
			String line;
			int lineNum = 0;
			System.out.println("Training Begin");
			while ((line = br.readLine()) != null) {
				if (lineNum % 1000000 == 0)
					System.out.println(lineNum);
				lineNum++;
				String[] paras = line.split("\t");// id seg1|seg2|seg3....
				double click = Double.parseDouble(paras[0]);
				double impression = Double.parseDouble(paras[1]);
				int queryID = Integer.parseInt(paras[7]);
				int purchaseID = Integer.parseInt(paras[8]);
				int titleID = Integer.parseInt(paras[9]);
				int descriptionID = Integer.parseInt(paras[10]);
				int userID = Integer.parseInt(paras[11]);

				// get all the tokens based on the hashmap created in
				// track2.data
				String query = getToken("Q", queryID);
				String title = getToken("T", titleID);
				String purchase = getToken("P", purchaseID);
				String description = getToken("D", descriptionID);
				String user = getToken("U", userID);

				// update the training data for each feature.
				totalClk += click;
				totalImp += impression;
				// should handle not found situation
				if (query != null)
					update("Q", query, click, impression);
				if (title != null)
					update("T", title, click, impression);
				if (purchase != null)
					update("P", purchase, click, impression);
				if (description != null)
					update("D", description, click, impression);
				if (user != null)
					update("U", user, click, impression);

			}
			br.close();
			System.out.println("Training finished");
		} catch (Exception e) {
			e.printStackTrace();
		}

	}

	// update function: after get the value from hashmap in track2.data, we need
	// to update the hashmap of each training feature.
	private static void update(String flag, String value, double click,
			double impression) {
		String[] sp = value.split("\\|");
		double nClick = impression - click;
		if (flag.equals("Q")) {// Query
			for (int i = 0; i < sp.length; i++) {
				updateMap(queryC, sp[i], click);
				updateMap(queryNC, sp[i], nClick);
			}
			queryCCount += click * sp.length;
			queryNCCount += nClick * sp.length;
		} else if (flag.equals("T")) {// Title
			for (int i = 0; i < sp.length; i++) {
				updateMap(titleC, sp[i], click);
				updateMap(titleNC, sp[i], nClick);
			}
			titleCCount += click * sp.length;
			titleNCCount += nClick * sp.length;
		} else if (flag.equals("P")) {// Purchase
			for (int i = 0; i < sp.length; i++) {
				updateMap(purchaseC, sp[i], click);
				updateMap(purchaseNC, sp[i], nClick);
			}
			purchaseCCount += click * sp.length;
			purchaseNCCount += nClick * sp.length;
		} else if (flag.equals("D")) {// Description
			for (int i = 0; i < sp.length; i++) {
				updateMap(descriptionC, sp[i], click);
				updateMap(descriptionNC, sp[i], nClick);
			}
			desCCount += click * sp.length;
			desNCCount += nClick * sp.length;
		}
		// User is different since sp[0] is gender and sp[1] is age
		if (flag.equals("U")) {// User
			genderCCount += click;
			genderNCCount += nClick;
			ageCCount += click;
			ageNCCount += nClick;
			updateMap(genderC, sp[0], click);
			updateMap(genderNC, sp[0], nClick);
			updateMap(ageC, sp[1], click);
			updateMap(ageNC, sp[1], nClick);
		}

	}

	private static void updateMap(HashMap<String, Double> map, String key,
			double addTimes) {
		if (map.containsKey(key))
			map.put(key, map.get(key) + addTimes);
		else if (!map.containsKey(key)) {
			map.put(key, addTimes);
		}

	}

	// get the tokenId from HashMap in track2.data
	/*
	 * to get a map: TitleMap = TitleToken.getInstance().getMap()
	 */
	private static String getToken(String flag, int id) {
		String result = null;
		if (flag.equals("Q")) {// QueryID
			result = QueryToken.getInstance().getMap().get(id);
		} else if (flag.equals("T")) {// TitleID
			result = TitleToken.getInstance().getMap().get(id);
		} else if (flag.equals("P")) {// PurchaseID
			result = PurchaseKeywordToken.getInstance().getMap().get(id);
		} else if (flag.equals("D")) {// DescriptionID
			result = DescriptionToken.getInstance().getMap().get(id);
		} else if (flag.equals("U")) {// UserID
			result = UserProfile.getInstance().getMap().get(id);
		}
		return result;
	}

	private static void computeProbability(NaiveBayesFeatures bayes, String flag) {
		if (flag.equals("Q")) {
//			for (int i = 0; i < bayes.queryClick.length; i++) {
//				bayes.queryClick[i] = 1.0 / (queryCCount + Config.tokenSize);
//				bayes.queryNotClick[i] = 1.0 / (queryNCCount + Config.tokenSize);
//			}
			compute(bayes.queryClick, queryC, queryCCount, Config.tokenSize);
			compute(bayes.queryNotClick, queryNC, queryNCCount,
					Config.tokenSize);
		} else if (flag.equals("T")) {
//			for (int i = 0; i < bayes.titleClick.length; i++) {
//				bayes.titleClick[i] = 1.0 / (titleCCount + Config.tokenSize);
//				bayes.titleNotClick[i] = 1.0 / (titleNCCount + Config.tokenSize);
//			}
			compute(bayes.titleClick, titleC, titleCCount, Config.tokenSize);
			compute(bayes.titleNotClick, titleNC, titleNCCount,
					Config.tokenSize);
		} else if (flag.equals("P")) {
//			for (int i = 0; i < bayes.purchaseClick.length; i++) {
//				bayes.purchaseClick[i] = 1.0 / (purchaseCCount + Config.tokenSize);
//				bayes.purchaseNotClick[i] = 1.0 / (purchaseNCCount + Config.tokenSize);
//			}
			compute(bayes.purchaseClick, purchaseC, purchaseCCount,
					Config.tokenSize);
			compute(bayes.purchaseNotClick, purchaseNC, purchaseNCCount,
					Config.tokenSize);
		} else if (flag.equals("D")) {
//			for (int i = 0; i < bayes.descriptionClick.length; i++) {
//				bayes.descriptionClick[i] = 1.0 / (desCCount + Config.tokenSize);
//				bayes.descriptionNotClick[i] = 1.0 / (desNCCount + Config.tokenSize);
//			}
			compute(bayes.descriptionClick, descriptionC, desCCount,
					Config.tokenSize);
			compute(bayes.descriptionNotClick, descriptionNC, desNCCount,
					Config.tokenSize);
		} else if (flag.equals("U")) {// UserID
			compute(bayes.genderClick, genderC, genderCCount, 3.0);
			compute(bayes.genderNotClick, genderNC, genderNCCount, 3.0);
			compute(bayes.ageClick, ageC, ageCCount, 7.0);
			compute(bayes.ageNotClick, ageNC, ageNCCount, 7.0);
		} else if (flag.equals("Prior")) {
			bayes.click = (totalClk + 1.0) / (totalImp + 2.0);
		}
	}

	// compute probability using formula and store in NaiveBayesFeatures
	private static void compute(double[] prob, HashMap<String, Double> mapC,
			double count, double size) {
		java.util.Iterator<Entry<String, Double>> iter = mapC.entrySet()
				.iterator();
		while (iter.hasNext()) {
			Entry<String, Double> entry = iter.next();
			String key = entry.getKey();
			double value = entry.getValue();
			double probability = (value + 1) / (count + size * 1.0);
			int index = Integer.parseInt(key);
			prob[index] = probability;
		}

	}

}
